class Actor(nn.Module):  
    def __init__(self, input_dim=None, node_num=node_num):  
        super(Actor, self).__init__()  

        # 图神经网络编码器  
        self.graph_encoder = GraphEncoder(  
            user_dim=3,  # 用户信息维度  
            node_dim=14,  # 节点信息维度  
            hidden_dim=64,  # 隐藏层维度  
        )  

        # 计算输入维度  
        def calculate_input_dim():  
            # 用户图特征  
            user_features_dim = user_num * 64  
            # 节点图特征  
            node_features_dim = node_num * 64  
            # 节点层和任务层特征  
            additional_features_dim = node_num * 2 * layer_num + 4 + 2 * layer_num  
            return user_features_dim + node_features_dim + additional_features_dim  

        # 全连接网络  
        self.fc1 = nn.Linear(calculate_input_dim(), 128)  
        self.fc2 = nn.Linear(128, 64)  
        self.fc3 = nn.Linear(64, 16)  
        self.action_head = nn.Linear(16, node_num + 1)  # +1 for cloud  

    def forward(self, x):  
        user_dim = 3 * user_num  
        node_dim = 14 * node_num  
        nodes_layer_dim = 2 * layer_num  # 2 * 719  
        tasks_dim = 4  
        tasks_layer_dim = layer_num  # 719  

        users = x[:, :user_dim].view(-1, user_num, 3)  
        nodes = x[:, user_dim : user_dim + node_dim].view(-1, node_num, 14)  

        # 正确读取 nodes_layers: (9, 2, 719)  
        nodes_layers = x[  
            :, user_dim + node_dim : user_dim + node_dim + nodes_layer_dim * node_num  
        ].view(-1, node_num, 2, layer_num)  

        tasks = x[  
            :,  
            user_dim  
            + node_dim  
            + nodes_layer_dim * node_num : user_dim  
            + node_dim  
            + nodes_layer_dim * node_num  
            + tasks_dim,  
        ]  

        # 正确读取 tasks_layers: (2, 719)  
        tasks_layers = x[  
            :, user_dim + node_dim + nodes_layer_dim * node_num + tasks_dim :  
        ].view(-1, 2, layer_num)  

        # 图编码  
        user_graph_features, node_graph_features = self.graph_encoder(users, nodes)  

        # 合并所有特征  
        combined_features = torch.cat(  
            [  
                user_graph_features.view(x.size(0), -1),  
                node_graph_features.view(x.size(0), -1),  
                nodes_layers.view(x.size(0), -1),  
                tasks,  
                tasks_layers.view(x.size(0), -1)  
            ],  
            dim=1,  
        )  

        # 全连接网络处理  
        x = F.leaky_relu(self.fc1(combined_features))  
        x = F.leaky_relu(self.fc2(x))  
        x = F.leaky_relu(self.fc3(x))  
        x = self.action_head(x)  
        action_prob = F.softmax(x, dim=1)  
        return action_prob  


class Critic(nn.Module):  
    def __init__(self, input_dim=None):  
        super(Critic, self).__init__()  

        # 图神经网络编码器  
        self.graph_encoder = GraphEncoder(  
            user_dim=3,  # 用户信息维度  
            node_dim=14,  # 节点信息维度  
            hidden_dim=64,  # 隐藏层维度  
        )  

        # 计算输入维度  
        def calculate_input_dim():  
            # 用户图特征  
            user_features_dim = user_num * 64  
            # 节点图特征  
            node_features_dim = node_num * 64  
            # 节点层和任务层特征  
            additional_features_dim = node_num * 2 * layer_num + 4 + 2 * layer_num  
            return user_features_dim + node_features_dim + additional_features_dim  

        # 全连接网络  
        self.fc1 = nn.Linear(calculate_input_dim(), 128)  
        self.fc2 = nn.Linear(128, 64)  
        self.fc3 = nn.Linear(64, 16)  
        self.value_head = nn.Linear(16, 1)  

    def forward(self, x):  
        user_dim = 3 * user_num  
        node_dim = 14 * node_num  
        nodes_layer_dim = 2 * layer_num  # 2 * 719  
        tasks_dim = 4  
        tasks_layer_dim = layer_num  # 719  

        users = x[:, :user_dim].view(-1, user_num, 3)  
        nodes = x[:, user_dim : user_dim + node_dim].view(-1, node_num, 14)  

        # 正确读取 nodes_layers: (9, 2, 719)  
        nodes_layers = x[  
            :, user_dim + node_dim : user_dim + node_dim + nodes_layer_dim * node_num  
        ].view(-1, node_num, 2, layer_num)  

        tasks = x[  
            :,  
            user_dim  
            + node_dim  
            + nodes_layer_dim * node_num : user_dim  
            + node_dim  
            + nodes_layer_dim * node_num  
            + tasks_dim,  
        ]  

        # 正确读取 tasks_layers: (2, 719)  
        tasks_layers = x[  
            :, user_dim + node_dim + nodes_layer_dim * node_num + tasks_dim :  
        ].view(-1, 2, layer_num)  

        # 图编码  
        user_graph_features, node_graph_features = self.graph_encoder(users, nodes)  

        # 合并所有特征  
        combined_features = torch.cat(  
            [  
                user_graph_features.view(x.size(0), -1),  
                node_graph_features.view(x.size(0), -1),  
                nodes_layers.view(x.size(0), -1),  
                tasks,  
                tasks_layers.view(x.size(0), -1)  
            ],  
            dim=1,  
        )  

        # 全连接网络处理  
        x = F.leaky_relu(self.fc1(combined_features))  
        x = F.leaky_relu(self.fc2(x))  
        x = F.leaky_relu(self.fc3(x))  
        value = self.value_head(x)  
        return value  